#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/5/25 10:47
@desc: remove the background
"""
from pathlib import Path

import h5py
import numpy as np
from skimage import measure, segmentation

from utils.config import DataPipe as DP
from utils.config import Number
from utils.data import Image
from utils.itk_fun import re_sample
from utils.log import get_console_logger


class BackGround(object):
    __slots__ = ('msg',
                 'origin',
                 'sampled',
                 'mask',
                 'bbox',
                 'sampled_size',
                 'save',
                 'logger')

    def __init__(self, origin: Image, save: str=None, logger=None):
        self.msg = None
        self.origin = origin
        self.save = save
        if logger is not None:
            self.logger = logger
        else:
            self.logger = get_console_logger('BackGround')

    def _get_mask_and_bounding_box(self,
                                   image: Image,
                                   mask_needed: bool=False,
                                   min_pixels: float=1e5) -> (np.ndarray,
                                                              np.ndarray):
        threshold = Image(image)
        shape = image.image_array.shape
        threshold.image_array = (
            np.less(image.image_array, Number.THRESHOLD))
        image_array = threshold.image_array.copy()
        for i in range(image_array.shape[0]):
            image_array[i, :, :] = (
                segmentation.clear_border(image_array[i, :, :]))
        regions = measure.regionprops(
            measure.label(image_array, connectivity=2))
        big_regions = []
        for region in regions:
            if region.area > min_pixels:
                big_regions.append(region)
        for i in range(image_array.shape[2]):
            image_array[:, :, i] = (
                segmentation.clear_border(image_array[:, :, i]))

        boxes = []
        if mask_needed:
            mask = np.zeros(shape, dtype=np.bool)
        else:
            mask = np.zeros([1, 1], dtype=np.bool)
        for region in big_regions:
            box = region.bbox
            if np.sum(image_array[box[0]:box[3],
                      box[1]:box[4],
                      box[2]:box[5]]) > 0:
                boxes.append(box)
                if mask_needed:
                    mask[np.split(region.coords, 3, 1)] = 1
        if len(boxes) == 0:
            self.msg = 'No Lung Find !'
            return None, None

        boxes = np.asarray(boxes)
        min_axis, max_axis = np.split(boxes, 2, 1)
        out_bbox = np.concatenate([np.min(min_axis, axis=0),
                                   np.max(max_axis, axis=0)], axis=0)
        return mask, out_bbox

    def _re_sample(self):
        image = self.origin
        _, bbox = self._get_mask_and_bounding_box(image)
        if bbox is None:
            return
        self.bbox = bbox
        shape = image.image_array.shape
        min_axis, max_axis = np.split(bbox, 2, 0)
        t_bbox = np.concatenate([np.maximum(min_axis - Number.ADD_SIZE, 0),
                                 np.minimum(max_axis + Number.ADD_SIZE,
                                            np.asarray(shape))],
                                axis=0)
        sub_origin = Image(image)
        sub_origin.origin = [t_bbox[2] * image.spacing[0],
                             t_bbox[1] * image.spacing[1],
                             t_bbox[0] * image.spacing[2]]
        sub_origin.image_array = (
            image.image_array[t_bbox[0]:t_bbox[3],
                              t_bbox[1]:t_bbox[4],
                              t_bbox[2]:t_bbox[5]])
        sampled = Image(image)
        sampled.init_with_itk(
            re_sample(sub_origin.to_itk(), Number.RESAMPLE_SPACING))
        self.sampled_size = sampled.image_array.shape
        self.sampled = sampled

    def _lung_mask(self):
        sampled = self.sampled
        mask = Image(sampled)
        mask_array, _ = self._get_mask_and_bounding_box(sampled, True, 1e6)
        if mask_array is None:
            return
        mask.image_array = mask_array
        self.mask = mask

    def _storage(self):
        path = Path(self.save) / DP.BACK + '.h5'
        with h5py.File(str(path), 'w') as file:
            file.create_dataset(DP.SAMPLED, data=self.sampled.image_array,
                                compression='gzip', compression_opts=9)
            file.create_dataset(DP.LUNG_MASK, data=self.mask.image_array,
                                compression='gzip', compression_opts=9)

    def update(self):
        self.logger.info('Removing the BackGround ...')
        self.logger.info('Re-sampling the lung')
        self._re_sample()
        self.logger.info('Removing the background and Getting the lung mask')
        self._lung_mask()
        if self.save is not None:
            self.logger.info('Saving the data')
            self._storage()
